Custom Task Development Guide

Overview

This guide will help you understand and create custom tasks in the LW-BenchHub. We'll analyze existing tasks and show you how to create your own step by step.

Quick Start

Before diving deep, here's what you need to know:

What is a Task? A task in LW-BenchHub defines:

• What the robot should accomplish (e.g., "open a drawer")
• Where objects are placed in the scene
• How success is measured
• What observations the robot receives (RL related)
• How rewards are calculated (RL related)
• Other events (RL related - domain randomization, etc.)

Key Components:

1. Task Configuration: Inherits from LwTaskBase
2. Scene Setup: Defines kitchen fixtures and robot placement
3. Success Criteria: Determines when the task is complete
4. Object Placement: Configures where items appear in the scene
5. Robot Placement Calculation: Decide where to spawn the robot so it can best complete the task.
6. Language Description: Human-readable task description
7. Other Configuration for RL： Observations, rewards, events, etc

Part I: Understanding Existing Tasks

Let's examine the CheesyBread task to understand the structure:

1. Basic Task Structure

# File: tasks/single_stage/kitchen_drawer.py
class CheesyBread(LwTaskBase):
    """Base class for drawer manipulation tasks."""
    
    # Task parameters
    layout_registry_names: list[int] = [FixtureType.COUNTER_NON_CORNER]
    task_name: str = "CheesyBread"

    def _setup_kitchen_references(self, scene):
        super()._setup_kitchen_references(scene)
        pass
    
    def _setup_scene(self, env, env_ids):
        super()._setup_scene(env, env_ids)
        pass
    
    def get_ep_meta(self):
        ep_meta = super().get_ep_meta()
        return ep_meta

    def _get_obj_cfgs(self):
        return []
    
    def _check_success(self, env):
        return torch.tensor([False], device=env.device).repeat(env.num_envs)

Key inheritance:

• LwTaskBase: Provides kitchen environment setup and basic task infrastructure (observations, rewards, etc.)

2. Core Methods Every Task Needs

A. Setup Kitchen References

def _setup_kitchen_references(self, scene):
    """Find and register kitchen fixtures for the task."""
    super()._setup_kitchen_references(scene)
    
    self.counter = self.register_fixture_ref(
        "counter", dict(id=FixtureType.COUNTER_NON_CORNER, size=(0.6, 0.6))
    )
    self.init_robot_base_ref = self.counter

B. Scene Initialization

def _setup_scene(self, env, env_ids):
    """Set initial scene state based on task requirements."""
    super()._setup_scene(env_ids)

C. Success Criteria

def _check_success(self):
    """Determine if the task has been completed successfully."""
    return OU.check_obj_in_receptacle(env, "bread", "bread_container") &\
        OU.gripper_obj_far(env, obj_name="cheese") &\
        OU.check_contact(env, self.objects["cheese"], self.objects["bread"])

D. Language Description

def get_ep_meta(self):
    """Generate human-readable task description."""
    ep_meta = super().get_ep_meta()
    ep_meta["lang"] = "Pick up the wedge of cheese and place it on the slice of bread to prepare a simple cheese on bread dish."
    return ep_meta

3. Smart Robot Placement

By setting self.init_robot_base_ref = self.counter, the framework will automatically compute a suitable initial robot position near the specified fixture (here, the counter), and place the robot into the scene accordingly. You do not need to manually specify the robot’s coordinates.

4. Default Configuration System

The LwTaskBase provides default configurations:

Observations (RL related)

class TaskBasePolicyObservationCfg(ObsGroup):
    """Observations for policy group."""

    def __post_init__(self):
        self.enable_corruption = True
        self.concatenate_terms = False

Rewards (RL related)

class RewardsCfg:
    action_rate_l2 = RewTerm(func=mdp.action_rate_l2, weight=-1e-2)  # Smooth actions
    joint_vel = RewTerm(func=mdp.joint_vel_l2, weight=-0.0001)       # Penalize fast movements

Events - Domain Randomization (RL related)

class EventCfg:
    init_task: EventTerm = MISSING
    reset_all = EventTerm(func=mdp.reset_scene_to_default, mode="reset")

Part II: Creating Your Custom Task

Let's create a WashDishes task step by step:

Step 1: Create the Task File

Create lw_benchhub_tasks/lightwheel_robocasa_tasks/single_stage/wash_dishes.py:

import torch
from lw_benchhub.core.models.fixtures import FixtureType
from lw_benchhub.core.tasks.base import LwTaskBase
import lw_benchhub.utils.object_utils as OU

class WashDishes(LwTaskBase):
    """Custom task: Move dirty dishes from counter to sink."""
    
    task_name: str = "WashDishes"

    def _setup_kitchen_references(self, scene):
        """Setup kitchen fixtures needed for the task."""
        super()._setup_kitchen_references(scene)
        
        # Register sink as the target location
        self.sink = self.register_fixture_ref("sink", dict(id=FixtureType.SINK))
        
        # Register counter near the sink for dish placement
        self.counter = self.register_fixture_ref(
            "counter", dict(id=FixtureType.COUNTER, ref=self.sink)
        )
        
        # Place robot near the sink
        self.init_robot_base_ref = self.sink

    def get_ep_meta(self):
        """Generate task description."""
        ep_meta = super().get_ep_meta()
        ep_meta["lang"] = "Pick up the dirty dish from the counter and place it in the sink."
        return ep_meta

    def _get_obj_cfgs(self):
        """Configure object placement in the scene."""
        cfgs = []

        # Main task object: dirty dish on counter
        cfgs.append(
            dict(
                name="dirty_dish",
                obj_groups="dishes",  # Use dish objects
                graspable=True,
                placement=dict(
                    fixture=self.counter,
                    sample_region_kwargs=dict(ref=self.sink),
                    size=(0.50, 0.30),           # Placement area size
                    pos=(None, -0.8),            # Position relative to counter
                    offset=(0.0, 0.10),          # Small offset from edge
                ),
            )
        )

        # Add some distractor objects to make it more realistic
        num_distractors = self.rng.integers(1, 3)
        for i in range(num_distractors):
            cfgs.append(
                dict(
                    name=f"clean_dish_{i+1}",
                    obj_groups="dishes",
                    placement=dict(
                        fixture=self.counter,
                        sample_region_kwargs=dict(ref=self.sink),
                        size=(0.80, 0.40),
                        pos=(None, -1.0),
                        offset=(0.0, 0.15),
                    ),
                )
            )

        return cfgs

    def _check_success(self, env):
        """Check if dish has been successfully moved to sink."""
        # Check if the main dish is inside the sink
        dish_in_sink = OU.obj_inside_of(env, "dirty_dish", self.sink)
        
        # Check if robot has released the dish (not grasping)
        gripper_released = OU.gripper_obj_far(env)
        
        return dish_in_sink & gripper_released

Step 2: Register Your Task

Add to lw_benchhub_tasks/lightwheel_robocasa_tasks/single_stage/__init__.py:

gym.register(
    id="Robocasa-Task-WashDishes",
    entry_point="isaaclab.envs:ManagerBasedRLEnv",
    kwargs={
        "env_cfg_entry_point": f"{__name__}.wash_dishes:WashDishes",
    },
    disable_env_checker=True,
)

Step 3: Test Your Task

import gymnasium as gym
from lw_benchhub.utils.env import parse_env_cfg, ExecuteMode

# Create and test your task
env_cfg = parse_env_cfg(
    scene_backend="robocasa",
    task_backend="robocasa",
    scene_name="robocasakitchen-1-1",
    robot_name="G1-Controller",
    task_name="WashDishes",
    robot_scale=1.0,
    execute_mode=ExecuteMode.TELEOP,
)
gym.register(
    id="test_WashDishes",
    entry_point="isaaclab.envs:ManagerBasedRLEnv",
    kwargs={},
    disable_env_checker=True,
)
env = gym.make("test_WashDishes").unwrapped
obs, info = env.reset()

# Run a few steps
for i in range(100):
    action = env.action_space.sample()  # Random action for testing
    obs, reward, terminated, truncated, info = env.step(action)
    
    if terminated:
        print(f"Task completed at step {i}!")
        break

Part III: Advanced Customization

While both task and robot definitions can individually specify their RL configurations (such as rewards, events, and observations), an effective RL task requires explicit binding between the task and the robot. This is essential because not all robots share the same joint or sensor names, and different tasks may require distinct actions. LW-BenchHub introduces a dedicated RL configuration layer to connect tasks with robots, consolidating and specifying observations, rewards, and event handling for each unique task-robot combination.

1) Where to add files (lw_benchhub_rl)

Create a task folder under lw_benchhub_rl/:

lw_benchhub_rl/
  my_task/
    __init__.py
    my_task_rl.py          # RL env cfg for the task
    mdp/                # Optional: task-specific obs/reward helpers
      observations.py
      rewards.py
    agents/
      skrl_ppo_cfg.yaml # Agent config (or other algs)

Register your RL env and agent in lw_benchhub_rl/__init__.py using a helper:

import gymnasium as gym
from .my_task import agents as my_task_agents

gym.register(
    id="Robocasa-Rl-RobotTaskRL",
    entry_point="isaaclab.envs:ManagerBasedRLEnv",
    kwargs={
        "env_cfg_entry_point": f"{__name__}.my_task.my_task_rl:RobotTaskRL",
        "skrl_cfg_entry_point": f"{my_task_agents.__name__}:skrl_ppo_cfg.yaml"
        "rsl_rl_cfg_entry_point": f"{my_task_agents.__name__}:rsl_rl_ppo_cfg:PPORunnerCfg"
    },
    disable_env_checker=True,
)

2) Robot-Task RL variant (robot-task dependent)

from lw_benchhub.core.rl.base import LwRL
from lw_benchhub.utils.decorators import rl_on

@rl_on(task=task_A)
@rl_on(task=task_A)
@rl_on(embodiment=robot_A)
@rl_on(embodiment=robot_B)
class RobotTaskRL(LwRL):
    def __init__(self):
        super().__init__()
        # define or reuse cfg here
        self.rewards_cfg = RewardsCfg()
        self.events_cfg = EventCfg()
        self.curriculum_cfg = CurriculumCfg()
        self.commands_cfg = CommandsCfg()
        self.policy_observation_cfg = PolicyObsCfg()
        self.resample_objects_placement_on_reset = False
        self.resample_robot_placement_on_reset = False
    
    def setup_env_config(self, orchestrator):
        super().setup_env_config(orchestrator)
        # Specify the precise body name for robot interaction in the task
        orchestrator.task.commands_cfg.object_pose.body_name = "right_wrist_yaw_link"
    
    def modify_env_cfg(self, env_cfg: IsaacLabArenaManagerBasedRLEnvCfg):
        super().modify_env_cfg(env_cfg)
        # Post-process the composed env_cfg (after composition of IsaacLab-Arena)
        self.setup_camera_and_foreground(env_cfg)

Key notes

• The @rl_on decorator associates this RL configuration with specific robots and tasks. Notably, a single configuration can be linked to multiple robots and multiple tasks, enabling shared handling logic across a group of robot-task pairs. For example, a configuration named DoubleArmPnPRL can be registered to serve various robots such as DoublePiper, DoublePanda, and X7s, as well as tasks like PnPCounterToCabinet, PnPSinkToCounter, and PnPStoveToCounter. This flexible assignment allows the RL configuration to support all combinations within the specified sets.
• Upon initialization (when the RL configuration is referenced via the rl field in the launch configuration), the system verifies that the active robot and task are included within the configuration’s supported domains. If not, an error is raised, thus ensuring the validity and consistency of the RL environment logic.
• For detailed examples, see lw_benchhub_rl/lift_obj/lift_obj.py.

Policy / General Observations

In observation design, we distinguish between general (non-policy) observations and policy-specific observations. This separation allows for seamless replacement of policy observations in RL configs without modifying the common observations, ensuring a clear and non-interfering scope for each type.

from lw_benchhub.core.rl.base import RlBasePolicyObservationCfg
import lw_benchhub.core.mdp as mdp
from lw_benchhub.core.rl.base import LwRL
from lw_benchhub.utils.decorators import rl_on
from isaaclab.managers import ObservationTermCfg as ObsTerm

@configclass
class RobotTaskPolicyObsCfg(RlBasePolicyObservationCfg):
    obj_pos = ObsTerm(func=mdp.funcA)

@rl_on(task=task_A)
@rl_on(task=task_A)
@rl_on(embodiment=robot_A)
@rl_on(embodiment=robot_B)
class RobotTaskRL(LwRL):
    def __init__(self):
        super().__init__()
        self.policy_observation_cfg = RobotTaskPolicyObsCfg()

from lw_benchhub.core.robots.robot_arena_base import LwEmbodimentBase, EmbodimentBaseObservationCfg
import lw_benchhub.core.mdp as mdp
from isaaclab.managers import ObservationTermCfg as ObsTerm

@configclass
class RobotObservationsCfg(EmbodimentBaseObservationCfg):
    """Observation specifications for the MDP."""

    @configclass
    class G1GeneralObsCfg(EmbodimentGeneralObsCfg):
        gripper_pos: ObsTerm = ObsTerm(func=lw_benchhub_mdp.gripper_pos)

    embodiment_general_obs: G1GeneralObsCfg = G1GeneralObsCfg()

class RobotCfg(LwEmbodimentBase):

    def __init__(self, enable_cameras: bool = False, initial_pose: Pose | None = None):
        super().__init__(enable_cameras, initial_pose)
        self.name = "AnyRobot"
        self.observation_config = RobotObservationsCfg()

4) Agent config (SKRL)

Place under lw_benchhub_rl/my_task/agents/:

# lw_benchhub_rl/my_task/agents/skrl_ppo_cfg.yaml
agent:
  class: PPO
  rollouts: 32
  learning_epochs: 5
  mini_batches: 4
  learning_rate: 3.0e-04
trainer:
  class: SequentialTrainer
  timesteps: 80000

5) Training configuration and run

Global training YAML under configs/rl/skrl/ (maps robot + task + scene), e.g.:

# configs/rl/skrl/my_robot_my_task.yaml
_base_: rl_base
task: MyTask
robot: MyRobot-RL
layout: robocasakitchen-1-8
num_envs: 128
enable_cameras: false
variant: State

Then run:

python lw_benchhub/scripts/rl/train.py --task_config configs/rl/skrl/my_robot_my_task.yaml --headless

Evaluate:

python lw_benchhub/scripts/rl/play.py --task_config configs/rl/skrl/my_robot_my_task.yaml

Best Practices

1. Start Simple

Begin with basic functionality:

• Define clear success criteria
• Use simple object placement
• Test frequently during development

2. Use Existing Patterns

• Inherit from base classes when possible
• Follow naming conventions from existing tasks
• Leverage the built-in fixture and object systems

3. Design for Robustness

• Handle edge cases in success checking
• Use appropriate collision detection
• Test with different scene layouts

4. Make it Learnable (RL related)

• Provide dense reward signals
• Include relevant observations
• Balance task difficulty

Common Patterns

Multi-Object Tasks

def _get_obj_cfgs(self):
    cfgs = []
    
    # Multiple target objects
    for i in range(3):
        cfgs.append(dict(
            name=f"target_obj_{i}",
            obj_groups="dishes",
            graspable=True,
            placement=dict(fixture=self.counter, ...)
        ))
    
    return cfgs

Sequential Tasks

def _check_success(self, env):
    # Check multiple conditions in sequence
    dish_picked = OU.gripper_obj_close(env, "dirty_dish")
    dish_in_sink = OU.obj_inside_of(env, "dirty_dish", self.sink)
    water_turned_on = self.sink.get_state()["faucet_on"]
    
    return dish_picked & dish_in_sink & water_turned_on

Conditional Behaviors

def _setup_scene(self, env, env_ids=None):
    # Randomize initial conditions
    if self.rng.random() > 0.5:
        self.sink.turn_on_faucet(env=env)
    else:
        self.sink.turn_off_faucet(env=env)
    
    super()._setup_scene(env, env_ids)

Troubleshooting

Task not working?

1. Check that your task is registered correctly
2. Verify fixture references are valid
3. Ensure success criteria are reachable
4. Test object placement doesn't cause collisions

Robot placement issues?

1. Adjust placement offsets
2. Check collision detection logic
3. Verify fixture dimensions are correct

Poor learning performance? (RL related)

1. Review reward signal density
2. Check observation relevance
3. Ensure task difficulty is appropriate
4. Verify success criteria are clear

Summary

Creating custom tasks in LW-BenchHub involves:

1. Understanding the structure: Inherit from base classes and implement key methods
2. Defining the task: Set up fixtures, objects, and success criteria
3. Testing thoroughly: Verify functionality before training
4. Iterating: Refine rewards and observations based on learning performance

The framework provides powerful tools for creating diverse robotics tasks. Start with simple examples like WashDishes and gradually add complexity as needed.

Remember: good tasks have clear objectives, learnable reward structures, and robust success criteria. Happy coding!